import os
import re
import string
import tensorflow as tf
from tensorflow.keras import layers, losses

# 1. 数据路径
dataset_dir = r"D:\casket\keras_data\aclImdb"
train_dir = os.path.join(dataset_dir, 'train')
# r"D:\casket\keras_data\aclImdb\train"
test_dir = os.path.join(dataset_dir, 'test')
# r"D:\casket\keras_data\aclImdb\test"

# 2. 加载数据集（自动分训练/验证）
batch_size = 32
seed = 42  # 随机种子，42是一个彩蛋数值，这个值可以随便写比如114514，这是个伪随机，同样的数值运行的结果是一样的。
'''
text_dataset_from_directory()方法会寻找路径中的文件夹，
把文件夹中的子文件夹，按照系统自带的排序，把第文件夹中的所有子文件都打上标签，第一个0，第二个1，以此类推。
# 注意：只会寻找目录下的子文件夹，不会找文件。只会提取目录中的文件夹中的文件的文本（不论什么格式，都会尝试读取出文本），只会找一层。

text_dataset_from_directory 会自动把文件夹里的文件打上标签。
把这些文件读出来，组成一个 tf.data.Dataset 对象：
每个元素是 (text_tensor, label_tensor)
text 是字符串张量，label 是数字张量（0 或 1）
还会自动打包成 batch（这里 batch_size=32）
'''
# 获取训练集
raw_train_ds = tf.keras.utils.text_dataset_from_directory(
    train_dir,  # 要找的目录
    batch_size=batch_size,  # 每个 batch 的大小
    validation_split=0.2,  # 表示要把数据的 20% 作为验证集
    subset='training',  # 表示要拿训练集的数据
    seed=seed  # 随机种子
)
# 获取验证集
raw_val_ds = tf.keras.utils.text_dataset_from_directory(
    train_dir,
    batch_size=batch_size,
    validation_split=0.2,
    subset='validation',
    seed=seed
)
# 获取测试集
raw_test_ds = tf.keras.utils.text_dataset_from_directory(
    test_dir,
    batch_size=batch_size
)


# ^^^ 拿出了3个数据集，分别是训练集、验证集、测试集

# 3. 文本标准化 & 向量化层
def custom_standardization(input_data):
    lowercase = tf.strings.lower(input_data)  # 全部转成小写
    stripped_html = tf.strings.regex_replace(lowercase, '<br />', ' ')  # 去除html标签
    return tf.strings.regex_replace(stripped_html, '[%s]' % re.escape(string.punctuation), '')  # 去除标点符号


max_features = 10000  # 词表大小
sequence_length = 250  # 序列长度,短了补0,长了截断。

# 创建了一个 TextVectorization 层（预处理层），它会把文本标准化、向量化，并把词表限制为 max_features 个词
vectorize_layer = layers.TextVectorization(
    standardize=custom_standardization,  # 在分词前先用这个函数处理数据
    max_tokens=max_features,  # 限制词表大小，每次最多只保留 max_features 个词
    output_mode='int',  # 输出模式为整数
    output_sequence_length=sequence_length  # 输出序列长度为 sequence_length
)

# .map()方法：把数据集中的每一个元素都映射到一个函数上，这个函数会返回一个新的元素，然后把这个新的元素放到新的数据集中。
# 这里的映射函数是 vectorize_text()，它会把文本转成整数序列。
train_text = raw_train_ds.map(lambda x, y: x)  # 去掉训练集数据中的标签，只保留数据
vectorize_layer.adapt(train_text)  # 把要训练集数据传给层，让他学习出一个词表


# .adapt()后，层会扫一遍数据的所有出现过的单词，按照词的频率排序，挑出最常见的 max_features 个词，
# 每个词都会分配一个整数id，然后把文本中的每个词替换成对应的整数id。

# 这个函数会把文本转成整数序列，并把标签保留下来。
def vectorize_text(text, label):
    text = tf.expand_dims(text, -1)  # 给元素的最后一维增加一个维度，因为 keras 通常要求字符串是二维的
    return vectorize_layer(text), label


# 把数据传给vectorize_layer层，层会把文本按照词表转换成整数序列。
# 没有配对上序列的词会被替换成0，配对上的词会被替换成对应的整数id。

train_ds = raw_train_ds.map(vectorize_text)  # 把训练集映射成了一个包含整数序列和标签的新数据集
val_ds = raw_val_ds.map(vectorize_text)
test_ds = raw_test_ds.map(vectorize_text)

# 性能优化
AUTOTUNE = tf.data.AUTOTUNE  # 框架会根据硬件自动调整要取多少个bath合适。
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)
# .cahe()方法：缓存数据集，在内存中保留数据，避免重复读取。
# .prefetch()方法：预取数据集，在后台异步加载数据，提高数据集的处理速度。
# 提前准备好下一个bath数据，放在缓存中，避免重复读取。


# 搭建模型
embedding_dim = 16  # 词向量维度

model = tf.keras.Sequential([
    # 传入了：Your movie is just like a pile of dog feces.
    # 变成了：[1,2,3,4,5,6,7,8,9,10,0,0,0,0,0,0],传给了Embedding层。
    layers.Embedding(max_features + 1,  # 词表大小，要加一是因为要预留一个特殊的OOV编号，表示未知词
                     embedding_dim,  # 词向量维度
                     input_length=sequence_length  # 告诉层输入的句子有多少个词ID。这个参数不直接参与运算，是为了方便后续层推导形状
                     ),
    layers.GlobalAveragePooling1D(),  # 把数据平均化，得到一个固定长度的向量。
    # 这句话里所有词的含义加起来 ➜ 取个平均 ➜ 得到整句话的大概语义
    layers.Dense(1)
])

model.compile(
    # 使用二元交叉熵作为损失函数，因为标签只有 0 或 1 两种，from_logits=True 表示输入是 logits 形式
    loss=losses.BinaryCrossentropy(from_logits=True),
    # 如果真实值是 1，希望模型输出接近 1，如果真实值是 0，希望模型输出接近 0。
    # 最终模型输出的值是0~1之间，越大表示越符合标签1，越小表示越符合标签0。
    # 如果1表示差评的话，那模型输出的是差评的概率。
    optimizer='adam',
    # 用二分类的准确率（BinaryAccuracy）来作为额外的监控指标。
    metrics=[tf.metrics.BinaryAccuracy(threshold=0.0)]
    # 会把输出和阈值比较，判定是 0 还是 1。如果模型的输出层进行了 sigmoid 变换，那么threshold=0.5。没有则threshold=0.0
)

# 8. 训练
history = model.fit(train_ds, validation_data=val_ds, epochs=10)

# 9. 评估
loss, accuracy = model.evaluate(test_ds)
print("Test Loss:", loss)
print("Test Accuracy:", accuracy)
